package com.vivien.example.springbootdemo.service;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.multimodalconversation.MultiModalConversation;
import com.alibaba.dashscope.aigc.multimodalconversation.MultiModalConversationParam;
import com.alibaba.dashscope.aigc.multimodalconversation.MultiModalConversationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.MultiModalMessage;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.exception.UploadFileException;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.vivien.example.springbootdemo.dto.ChatMultiModalRequest;
import com.vivien.example.springbootdemo.enums.ChatPrompt;
import com.vivien.example.springbootdemo.exception.CustomException;
import com.vivien.example.springbootdemo.mapper.UserApiKeyMapper;
import com.vivien.example.springbootdemo.pojo.UserApiKey;
import io.reactivex.Flowable;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.time.Duration;
import java.util.*;

@Service
public class ChatService {
    public static final String API_KEY = "sk-4d8ba379e9c1415d967fc7d9785c6813";
    @Autowired
    private UserApiKeyMapper userApiKeyMapper;
    private static final Logger logger = LoggerFactory.getLogger(ChatService.class);
    private static StringBuilder fullContent = new StringBuilder();
    private final static List<Message> messages = new ArrayList<>();
    private final static List<MultiModalMessage> multiModalMessages = new ArrayList<>();

    /**
     * 大模型相关
     */
    private static void handleGenerationResult(GenerationResult message) {
        String content = message.getOutput().getChoices().get(0).getMessage().getContent();
        fullContent.append(content);
        System.out.println(content);
    }
    private static GenerationParam buildGenerationParam(List<Message> messages) {
        return GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(API_KEY)
                .model("qwen-plus")
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }
    private static MultiModalConversationParam buildGenerationParamMultiModal(List<MultiModalMessage> multiModalMessages) {
        List<MultiModalMessage> tempMessages = new ArrayList<>();
        for (MultiModalMessage multiModalMessage : multiModalMessages) {
            if(multiModalMessage.getContent().size() == 2) {
                List<Map<String, Object>> msgList = multiModalMessage.getContent();
                tempMessages.add(MultiModalMessage.builder().role(multiModalMessage.getRole()).content(Arrays.asList(
                        new HashMap<String, Object>(){{put("image", msgList.get(0).get("image"));}}, new HashMap<String, Object>(){{put("text", msgList.get(1).get("text"));}})).build());
            } else {
                tempMessages.add(MultiModalMessage.builder().role(multiModalMessage.getRole()).content(multiModalMessage.getContent()).build());
            }
        }
        return MultiModalConversationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(API_KEY)
                .model("qwen-vl-max-latest")
                .messages(tempMessages)
                .incrementalOutput(true)
                .build();
    }
    private static Message createMessage(Role role, String content) {
        return Message.builder().role(role.getValue()).content(content).build();
    }
    private static void logError(Exception e) {
        logger.error("An exception occurred: {}", e.getMessage());
    }
    public static String fullCallGenerationWithMessages(GenerationParam param) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        // 在此处清空fullContent
        fullContent = new StringBuilder();
        System.out.println("流式输出内容为：");
        Flowable<GenerationResult> result = gen.streamCall(param);
        result.blockingForEach(message -> handleGenerationResult(message));
        System.out.println("完整内容为: " + fullContent.toString());
        return fullContent.toString();
    }
    private static Flowable<GenerationResult> streamCallGenerationWithMessages(GenerationParam param) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        return gen.streamCall(param);
    }
    private static Flowable<MultiModalConversationResult> streamCallGenerationWithMessages(MultiModalConversationParam param) throws ApiException, NoApiKeyException, UploadFileException {
        MultiModalConversation conv = new MultiModalConversation();
        return conv.streamCall(param);
    }
    // 检查并设置初始化角色
    private void checkAndSetRole(String prompt) {
        if (messages.isEmpty()) {
            messages.add(createMessage(Role.SYSTEM, prompt));
        } else {
            messages.removeIf(message -> message.getContent() == null);
        }
    }
    private void checkAndSetRoleMultiModal(String prompt) {
        MultiModalMessage systemMessage = MultiModalMessage.builder().role(Role.SYSTEM.getValue())
                .content(Collections.singletonList(Collections.singletonMap("text", prompt))).build();
        multiModalMessages.add(systemMessage);
        multiModalMessages.removeIf(message -> message.getContent() == null);
    }
    private void addAnalyze(String prompt) {
        MultiModalMessage systemMessage = MultiModalMessage.builder().role(Role.SYSTEM.getValue())
                .content(Collections.singletonList(Collections.singletonMap("text", prompt))).build();
        multiModalMessages.add(systemMessage);
    }
    // 组合用户消息和API
    private GenerationParam buildUserMsg(String userInput) {
        messages.add(createMessage(Role.USER, userInput));
        return buildGenerationParam(messages);
    }
    // 组合用户消息和API
    private void addToMessages(String userInput, String img) {
        MultiModalMessage userMessage = img == null ?
                MultiModalMessage.builder().role(Role.USER.getValue())
                .content(Collections.singletonList(new HashMap<String, Object>() {{put("text", userInput);}})).build()
                :MultiModalMessage.builder().role(Role.USER.getValue())
                .content(Arrays.asList(new HashMap<String, Object>(){{put("image", img);}}, new HashMap<String, Object>(){{put("text", userInput);}})).build();
        multiModalMessages.add(userMessage);
    }
    private MultiModalConversationParam buildUserMsg(String userInput, String img) {
        addToMessages(userInput, img);
        return buildGenerationParamMultiModal(multiModalMessages);
    }
    // 解析保存图像
    public String decodeAndSaveImage(String name, String imgBase64) {
        // 指定缓存目录
        String cacheDir = System.getProperty("java.io.tmpdir"); // 系统临时目录
        File dir = new File(cacheDir);
        if (!dir.exists()) {
            dir.mkdirs(); // 创建目录
        }
        int n = 1; // 从 1 开始
        File imageFile;
        do {
            String fileName = name + "_" + n + ".jpg";
            imageFile = new File(dir, fileName);
            n++;
        } while (imageFile.exists()); // 如果文件已存在，n 递增
        int delimiterIndex = imgBase64.indexOf("|||");
        String base64Content = delimiterIndex != -1 ? imgBase64.substring(delimiterIndex + 3) : imgBase64;
        assert !base64Content.isEmpty() : "Base64 内容不能为空";
        try {
            byte[] imageBytes = Base64.getDecoder().decode(base64Content);
            try (FileOutputStream fos = new FileOutputStream(imageFile)) {
                fos.write(imageBytes);
            }
            System.out.println("图像保存到" + imageFile.getAbsolutePath());
            return imageFile.getAbsolutePath();
        } catch (IOException e) {
            throw new RuntimeException("保存图像失败: " + e.getMessage(), e);
        }
    }
    private String findApiKeyByName(String name) {
        QueryWrapper<UserApiKey> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("user_name", name).select("api_key");
        UserApiKey userApi = userApiKeyMapper.selectOne(queryWrapper);
        return userApi != null ? userApi.getApiKey() : null;
    }
    private String checkApi(String name) {
        return Objects.requireNonNull(findApiKeyByName(name),
                () -> {
                    throw new CustomException.NoApiKeyException("API key not found");
                });
    }
    public String textChat(String name, String userInput) {
        try {
            checkAndSetRole(ChatPrompt.INSURANCE_PROMPT.getMessage());
            return fullCallGenerationWithMessages(buildUserMsg(userInput));
        } catch (Exception  e) {
            logError(e);
        }
        return null;
    }

    private void analyzeInput(ChatMultiModalRequest chatMultiModalRequest) throws NoApiKeyException, UploadFileException {
        String name = chatMultiModalRequest.getName(),
                userInput = chatMultiModalRequest.getUserInput(),
                img = chatMultiModalRequest.getImg();
        if (img != null) {
            String imgPath = "file:///" + decodeAndSaveImage(name, img).replace("\\","/");
            Flowable<MultiModalConversationResult> result = streamCallGenerationWithMessages(buildUserMsg(userInput, imgPath));
            removeSysMsg();
            extractedFullMessage(result);
        } else {
            Flowable<MultiModalConversationResult> result = streamCallGenerationWithMessages(buildUserMsg(userInput, null));
            removeSysMsg();
            extractedFullMessage(result);
        }
        System.out.println("分析结果：" + fullContent);
        multiModalMessages.remove(multiModalMessages.size() - 1);
    }

    private static void removeSysMsg() {
        multiModalMessages.removeIf(message -> message.getRole().equals(Role.SYSTEM.getValue()));
    }

    public Flux<ServerSentEvent<String>> analyzeChat(ChatMultiModalRequest chatMultiModalRequest) throws NoApiKeyException, UploadFileException {
        // 检查apiKey
        checkApi(chatMultiModalRequest.getName());
        // 初始化角色
        addAnalyze(ChatPrompt.ANALYZE_PROMPT.getMessage());
        // 分析用户意图
        analyzeInput(chatMultiModalRequest);
        if(fullContent.toString().contains("推荐") || fullContent.toString().contains("咨询")) {
            return multiModalChat(chatMultiModalRequest, ChatPrompt.INSURANCE_PROMPT.getMessage());
        } else if (fullContent.toString().contains("理赔")) {
            return multiModalChat(chatMultiModalRequest, ChatPrompt.CLAIM_PROMPT.getMessage());
        } else if (fullContent.toString().contains("查询")) {
            return multiModalChat(chatMultiModalRequest, ChatPrompt.QUERY_PROMPT.getMessage());
        } else {
            return multiModalChat(chatMultiModalRequest, ChatPrompt.REJECT_PROMPT.getMessage());
        }
    }
    public Flux<ServerSentEvent<String>> multiModalChat(ChatMultiModalRequest chatMultiModalRequest, String prompt) {
        String name = chatMultiModalRequest.getName(),
                userInput = chatMultiModalRequest.getUserInput(),
                img = chatMultiModalRequest.getImg();
        try {
            checkAndSetRoleMultiModal(prompt);
            if (img != null) {
                String imgPath = "file:///" + decodeAndSaveImage(name, img).replace("\\","/");
                Flowable<MultiModalConversationResult> result = streamCallGenerationWithMessages(buildUserMsg(userInput, imgPath));
                removeSysMsg();
                return getServerSentEventFlux(result);
            } else {
                Flowable<MultiModalConversationResult> result = streamCallGenerationWithMessages(buildUserMsg(userInput, null));
                removeSysMsg();
                return getServerSentEventFlux(result);
            }
        } catch (Exception  e) {
            logError(e);
        }
        return null;
    }

    @NotNull
    private Flux<ServerSentEvent<String>> getServerSentEventFlux(Flowable<MultiModalConversationResult> result) {
        System.out.println("流式输出内容为：");
        extractedFullMessage(result);

        multiModalMessages.add(MultiModalMessage.builder().role(Role.ASSISTANT.getValue())
                .content(Collections.singletonList(new HashMap<String, Object>() {{put("text", fullContent.toString());}})).build());

        return Flux.from(result).delayElements(Duration.ofMillis(100))
                .map(message -> {
                    List<Map<String, Object>> content = message.getOutput().getChoices().get(0).getMessage().getContent();
                    String output = content.isEmpty() ? "" : content.get(0).get("text").toString();
                    System.out.println(output); // print the output
                    return ServerSentEvent.<String>builder()
                            .id(UUID.randomUUID().toString())
                            .event("message")
                            .data(output)
                            .build();
                });
    }

    private static void extractedFullMessage(Flowable<MultiModalConversationResult> result) {
        fullContent = new StringBuilder();
        result.blockingForEach(item -> {
            try {
                List<Map<String, Object>> content = item.getOutput().getChoices().get(0).getMessage().getContent();
                String output = content.isEmpty() ? "" : content.get(0).get("text").toString();
                fullContent.append(output);
                // System.out.println(JsonUtils.toJson(item));
            } catch (Exception e){
                System.exit(0);
            }
        });
    }

    //@Resource
    public Flux<ServerSentEvent<String>> streamChat(String name, String userInput) {
        try {
            checkAndSetRole(ChatPrompt.INSURANCE_PROMPT.getMessage());
            return Flux.from(streamCallGenerationWithMessages(buildUserMsg(userInput)))
                    // add delay between each event
                    .delayElements(Duration.ofMillis(1000))
                    .map(message -> {
                        String output = message.getOutput().getChoices().get(0).getMessage().getContent();
                        System.out.println(output); // print the output
                        return ServerSentEvent.<String>builder()
                                .id(UUID.randomUUID().toString())
                                .event("message")
                                .data(output)
                                .build();
                    });
        } catch (Exception  e) {
            logError(e);
        }
        return null;
    }

    public boolean translateMessage(String userInput){
        return false;
    }
}
